import os

import graphviz
import networkx as nx

from tools.path_utils import my_path_join


class GraphNode2Graphviz:
    def __init__(self, graph, out_folder, file_name, extention):
        self.graph: nx.DiGraph = graph
        my_graph_out_folder = my_path_join(out_folder)
        self.file_name = file_name
        self.extention = extention
        self.file_dot_name = f"{self.file_name}.dot"
        self.file_picture_name = f"{self.file_dot_name}.{extention}"
        self.file_dot_path = os.path.join(my_graph_out_folder, self.file_dot_name)
        self.file_picture_path = os.path.join(my_graph_out_folder, self.file_picture_name)

    def graph_node_2_dot(self):
        print(f"正在导出dot:{self.file_name}")
        # 使用DotExporter导出Dot格式的树结构
        with open(self.file_dot_path, "w", encoding='utf-8') as file:
            file.write("digraph {")
            file.write("\ngraph [bgcolor=lightblue rankdir=TB]")
            self.write_node_list(file)
            self.write_edge_list(file)
            file.write("\n}")

    def write_node_list(self, file):
        for node in self.graph.nodes():
            data = self.graph.nodes[node]['data']
            label, fillcolor = data.get_label()
            file.write(f"\n\"{data.timestamp}\"[label=\"{label}\" shape=box fillcolor={fillcolor} style=filled]")

    def write_edge_list(self, file):
        for edge in self.graph.edges():
            node1, node2 = edge

            data = self.graph.edges[edge]['data']
            label = data.get_label()

            file.write(f"\n\"{node1}\"->\"{node2}\"[label=\"{label}\"]")

    def graphviz_to_picture(self):
        print(f"正在绘制{self.file_name}")
        # 使用graphviz将Dot格式转换为图片
        graphviz.render('dot', self.extention, self.file_dot_path)

    def get_picture_file_path(self):
        return self.file_picture_path, self.file_picture_name, self.file_dot_path


def graph_node_2_graphviz_picture_main(graph, out_folder, out_file, extension="svg"):
    my_graph_node_2_graphviz = GraphNode2Graphviz(graph, out_folder, out_file, extension)
    my_graph_node_2_graphviz.graph_node_2_dot()
    my_graph_node_2_graphviz.graphviz_to_picture()
    file_picture_path, file_picture_name, file_dot_path = my_graph_node_2_graphviz.get_picture_file_path()
    return file_picture_path, file_picture_name, file_dot_path
